// Copyright 2022 Chen Jun

#ifndef ARMOR_DETECTOR__NUMBER_CLASSIFIER_HPP_
#define ARMOR_DETECTOR__NUMBER_CLASSIFIER_HPP_

// OpenCV
#include <opencv2/opencv.hpp>

// STL

#include <string>
#include <vector>

#include "hnurm_detect/armor.hpp"

namespace hnurm
{
class NumberClassifier
{
public:
    NumberClassifier(
        const std::string              &model_path,
        const std::string              &label_path,
        const double                    threshold,
        const std::vector<std::string> &ignore_classes
    );

    void extractNumbers(const cv::Mat &src, std::vector<Armor> &armors);

    void classify(std::vector<Armor> &armors,int self_color);

    double threshold;

private:
    cv::dnn::Net             net_;
    std::vector<std::string> class_names_;

    std::vector<std::string> ignore_classes_;
};
}  // namespace hnurm

#endif  // ARMOR_DETECTOR__NUMBER_CLASSIFIER_HPP_
